{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e462b448",
   "metadata": {},
   "source": [
    "# 08 JIT Control Flow and Logic\n",
    "\n",
    "Original Documentation: https://docs.jax.dev/en/latest/control-flow.html\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "adf06212",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d13223db",
   "metadata": {},
   "source": [
    "While executing eagerly outside of JIT, JAX works nicely with Python control flow.\n",
    "\n",
    "Since Python control flow over runtime values can dynamically change the size of the control flow graph, it is not possible for JIT to produce a compiled representation.\n",
    "\n",
    "Paths dependent on the dtype or shape are okay, but will require recompilation when a value of a new dtype or shape is provided.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2b948a06",
   "metadata": {},
   "outputs": [
    {
     "ename": "TracerBoolConversionError",
     "evalue": "Attempted boolean conversion of traced array with shape bool[].\nThe error occurred while tracing the function f at /var/folders/r0/zlgwr7551fg116vdpn_9s7ph0000gn/T/ipykernel_31847/3756118353.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument x.\nSee https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mTracerBoolConversionError\u001b[39m                 Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[2]\u001b[39m\u001b[32m, line 6\u001b[39m\n\u001b[32m      1\u001b[39m \u001b[38;5;129m@jax\u001b[39m.jit\n\u001b[32m      2\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mf\u001b[39m(x):\n\u001b[32m      3\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m (x > \u001b[32m0\u001b[39m) \u001b[38;5;129;01mand\u001b[39;00m (x < \u001b[32m3\u001b[39m)  \u001b[38;5;66;03m# This will fail! Runtime comparison in JIT.\u001b[39;00m\n\u001b[32m----> \u001b[39m\u001b[32m6\u001b[39m \u001b[43mf\u001b[49m\u001b[43m(\u001b[49m\u001b[32;43m1.0\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "    \u001b[31m[... skipping hidden 13 frame]\u001b[39m\n",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[2]\u001b[39m\u001b[32m, line 3\u001b[39m, in \u001b[36mf\u001b[39m\u001b[34m(x)\u001b[39m\n\u001b[32m      1\u001b[39m \u001b[38;5;129m@jax\u001b[39m.jit\n\u001b[32m      2\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mf\u001b[39m(x):\n\u001b[32m----> \u001b[39m\u001b[32m3\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m \u001b[49m\u001b[43m>\u001b[49m\u001b[43m \u001b[49m\u001b[32;43m0\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mand\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m \u001b[49m\u001b[43m<\u001b[49m\u001b[43m \u001b[49m\u001b[32;43m3\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "    \u001b[31m[... skipping hidden 1 frame]\u001b[39m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/Programming/Learn JAX/.venv/lib/python3.12/site-packages/jax/_src/core.py:1721\u001b[39m, in \u001b[36mconcretization_function_error.<locals>.error\u001b[39m\u001b[34m(self, arg)\u001b[39m\n\u001b[32m   1720\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34merror\u001b[39m(\u001b[38;5;28mself\u001b[39m, arg):\n\u001b[32m-> \u001b[39m\u001b[32m1721\u001b[39m   \u001b[38;5;28;01mraise\u001b[39;00m TracerBoolConversionError(arg)\n",
      "\u001b[31mTracerBoolConversionError\u001b[39m: Attempted boolean conversion of traced array with shape bool[].\nThe error occurred while tracing the function f at /var/folders/r0/zlgwr7551fg116vdpn_9s7ph0000gn/T/ipykernel_31847/3756118353.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument x.\nSee https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError"
     ]
    }
   ],
   "source": [
    "@jax.jit\n",
    "def f(x):\n",
    "    return (x > 0) and (x < 3)  # This will fail! Runtime comparison in JIT.\n",
    "\n",
    "\n",
    "f(1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d8918b8",
   "metadata": {},
   "source": [
    "When we JIT compile a function, we want to compile a version that works for many different values so we can cache and reuse the XLA without needing to recompile.\n",
    "\n",
    "The only way to still condition on a runtime value is to mark the argument as static (which treats the value as a constant, meaning new values passed to the static argument will require function recompilation).\n",
    "\n",
    "In a loop, if the loop range is constant or static, then JAX will unroll the loop.\n",
    "\n",
    "### Functions with argument-value dependent shapes\n",
    "\n",
    "A common case programmers forget to account for:\n",
    "\n",
    "Functions that we want to JIT compile cannot specialize the shapes of internal arrays on argument values either.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d61313f9",
   "metadata": {},
   "outputs": [
    {
     "ename": "ConcretizationTypeError",
     "evalue": "Abstract tracer value encountered where concrete value is expected: traced array with shape int32[]\nIt arose in the jnp.arange argument 'stop'\nThe error occurred while tracing the function alloc at /var/folders/r0/zlgwr7551fg116vdpn_9s7ph0000gn/T/ipykernel_31847/298980794.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument length.\n\nSee https://docs.jax.dev/en/latest/errors.html#jax.errors.ConcretizationTypeError",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mConcretizationTypeError\u001b[39m                   Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[3]\u001b[39m\u001b[32m, line 6\u001b[39m\n\u001b[32m      1\u001b[39m \u001b[38;5;129m@jax\u001b[39m.jit\n\u001b[32m      2\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34malloc\u001b[39m(length, val):\n\u001b[32m      3\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m jnp.arange(length) * val  \u001b[38;5;66;03m# Length of array can't be value dependent!\u001b[39;00m\n\u001b[32m----> \u001b[39m\u001b[32m6\u001b[39m \u001b[43malloc\u001b[49m\u001b[43m(\u001b[49m\u001b[32;43m10\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[32;43m1.0\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "    \u001b[31m[... skipping hidden 13 frame]\u001b[39m\n",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[3]\u001b[39m\u001b[32m, line 3\u001b[39m, in \u001b[36malloc\u001b[39m\u001b[34m(length, val)\u001b[39m\n\u001b[32m      1\u001b[39m \u001b[38;5;129m@jax\u001b[39m.jit\n\u001b[32m      2\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34malloc\u001b[39m(length, val):\n\u001b[32m----> \u001b[39m\u001b[32m3\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mjnp\u001b[49m\u001b[43m.\u001b[49m\u001b[43marange\u001b[49m\u001b[43m(\u001b[49m\u001b[43mlength\u001b[49m\u001b[43m)\u001b[49m * val\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/Programming/Learn JAX/.venv/lib/python3.12/site-packages/jax/_src/numpy/lax_numpy.py:5966\u001b[39m, in \u001b[36marange\u001b[39m\u001b[34m(start, stop, step, dtype, device, out_sharding)\u001b[39m\n\u001b[32m   5964\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m sharding \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m sharding._is_concrete:\n\u001b[32m   5965\u001b[39m   \u001b[38;5;28;01massert\u001b[39;00m sharding \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(sharding, NamedSharding)\n\u001b[32m-> \u001b[39m\u001b[32m5966\u001b[39m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_arange\u001b[49m\u001b[43m(\u001b[49m\u001b[43mstart\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstop\u001b[49m\u001b[43m=\u001b[49m\u001b[43mstop\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstep\u001b[49m\u001b[43m=\u001b[49m\u001b[43mstep\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdtype\u001b[49m\u001b[43m=\u001b[49m\u001b[43mdtype\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   5967\u001b[39m \u001b[43m                 \u001b[49m\u001b[43mout_sharding\u001b[49m\u001b[43m=\u001b[49m\u001b[43msharding\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   5968\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m   5969\u001b[39m   output = _arange(start, stop=stop, step=step, dtype=dtype)\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/Programming/Learn JAX/.venv/lib/python3.12/site-packages/jax/_src/numpy/lax_numpy.py:5980\u001b[39m, in \u001b[36m_arange\u001b[39m\u001b[34m(start, stop, step, dtype, out_sharding)\u001b[39m\n\u001b[32m   5978\u001b[39m util.check_arraylike(\u001b[33m\"\u001b[39m\u001b[33marange\u001b[39m\u001b[33m\"\u001b[39m, start)\n\u001b[32m   5979\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m stop \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m step \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[32m-> \u001b[39m\u001b[32m5980\u001b[39m   start = \u001b[43mcore\u001b[49m\u001b[43m.\u001b[49m\u001b[43mconcrete_or_error\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstart\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mIt arose in the jnp.arange argument \u001b[39;49m\u001b[33;43m'\u001b[39;49m\u001b[33;43mstop\u001b[39;49m\u001b[33;43m'\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[32m   5981\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m   5982\u001b[39m   start = core.concrete_or_error(\u001b[38;5;28;01mNone\u001b[39;00m, start, \u001b[33m\"\u001b[39m\u001b[33mIt arose in the jnp.arange argument \u001b[39m\u001b[33m'\u001b[39m\u001b[33mstart\u001b[39m\u001b[33m'\u001b[39m\u001b[33m\"\u001b[39m)\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/Programming/Learn JAX/.venv/lib/python3.12/site-packages/jax/_src/core.py:1737\u001b[39m, in \u001b[36mconcrete_or_error\u001b[39m\u001b[34m(force, val, context)\u001b[39m\n\u001b[32m   1735\u001b[39m maybe_concrete = val.to_concrete_value()\n\u001b[32m   1736\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m maybe_concrete \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[32m-> \u001b[39m\u001b[32m1737\u001b[39m   \u001b[38;5;28;01mraise\u001b[39;00m ConcretizationTypeError(val, context)\n\u001b[32m   1738\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m   1739\u001b[39m   \u001b[38;5;28;01mreturn\u001b[39;00m force(maybe_concrete)\n",
      "\u001b[31mConcretizationTypeError\u001b[39m: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[]\nIt arose in the jnp.arange argument 'stop'\nThe error occurred while tracing the function alloc at /var/folders/r0/zlgwr7551fg116vdpn_9s7ph0000gn/T/ipykernel_31847/298980794.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument length.\n\nSee https://docs.jax.dev/en/latest/errors.html#jax.errors.ConcretizationTypeError"
     ]
    }
   ],
   "source": [
    "@jax.jit\n",
    "def alloc(length, val):\n",
    "    return jnp.arange(length) * val  # Length of array can't be value dependent!\n",
    "\n",
    "\n",
    "alloc(10, 1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1533140f",
   "metadata": {},
   "source": [
    "## Structured control flow primitives\n",
    "\n",
    "To avoid requiring retracing and avoid unrolling loops in control flow conditioned on runtime values, JAX provides basic primitives.\n",
    "\n",
    "### Conditions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c74387ba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(11, dtype=int32, weak_type=True)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from jax import lax\n",
    "\n",
    "\n",
    "# Python version\n",
    "def cond(pred, true_fun, false_fun, operand):\n",
    "    if pred:\n",
    "        return true_fun(operand)\n",
    "    else:\n",
    "        return false_fun(operand)\n",
    "\n",
    "\n",
    "# JAX Version\n",
    "# lax.cond(pred, true_fun, false_fun, operand)\n",
    "lax.cond(True, lambda x: x + 1, lambda x: x - 1, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "419c799c",
   "metadata": {},
   "source": [
    "LAX also provides other functions for conditions:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8cab4f4b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(11, dtype=int32, weak_type=True)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 1) Choose between two precomputed values of the same shape and dtype\n",
    "# lax.select(pred, a, b)\n",
    "a = jnp.arange(10) * 10\n",
    "b = jnp.arange(10) * 5\n",
    "lax.select(a[-1] > b[-1], a, b)\n",
    "\n",
    "# 2) Like lax.cond, but pick between any number of branch functions\n",
    "# lax.switch(index, branches, operand)\n",
    "branches = [lambda x: x + 1, lambda x: x - 1, lambda x: x * 2]\n",
    "lax.switch(0, branches, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1e00e18",
   "metadata": {},
   "source": [
    "The `jax.numpy` API also provides conditional functions:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8fd72e9a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0, 2, 0, 2, 0, 1, 0, 2, 0, 2], dtype=int32, weak_type=True)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 1) Similar to lax.select but for numpy arrays\n",
    "# jnp.where(pred, a, b)\n",
    "a = jnp.arange(10)\n",
    "jnp.where(a % 2 == 0, a, 0)\n",
    "\n",
    "# 2) Similar to lax.switch but for numpy arrays\n",
    "# jnp.piecewise(a, conds, funcs)\n",
    "conds = [a % 2 == 0, a % 5 == 0]\n",
    "funcs = [lambda x: x - 1, lambda x: x + 1, lambda x: x]  # Last func is for `else` case\n",
    "jnp.piecewise(a, conds, funcs)\n",
    "\n",
    "# 3) jnp.select(conds, values)\n",
    "jnp.select(conds, [0, 1], default=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8020e034",
   "metadata": {},
   "source": [
    "### While Loops\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f18e45f6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(10, dtype=int32, weak_type=True)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Python version\n",
    "def while_loop(cond_fun, body_fun, init_val):\n",
    "    val = init_val\n",
    "    while cond_fun(val):\n",
    "        val = body_fun(val)\n",
    "    return val\n",
    "\n",
    "\n",
    "# JAX version\n",
    "# lax.while_loop(cond_fun, body_fun, init_val)\n",
    "lax.while_loop(lambda x: x < 10, lambda x: x + 1, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4109ea0",
   "metadata": {},
   "source": [
    "### For Loops\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c4386b73",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "45\n"
     ]
    }
   ],
   "source": [
    "# Python version\n",
    "def for_loop(start, stop, body_fun, init_val):\n",
    "    val = init_val\n",
    "    for i in range(start, stop):\n",
    "        val = body_fun(i, val)\n",
    "    return val\n",
    "\n",
    "\n",
    "# JAX version\n",
    "# lax.for_loop(start, stop, body_fun, init_val)\n",
    "print(lax.fori_loop(0, 10, lambda i, x: x + i, 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af75c5b8",
   "metadata": {},
   "source": [
    "## Logical Operators\n",
    "\n",
    "`jax.numpy` provides `logical_and`, `logical_or`, and `logical_not` which operate element-wise on arrays and can be ran in JIT contexts.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0f823748",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "def python_check_positive_even(x):\n",
    "    is_even = x % 2 == 0\n",
    "    return is_even and (x > 2)  # Short-circuits: when `is_even` is False, x > 0 is not evaluated\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def jax_check_positive_even(x):\n",
    "    is_even = x % 2 == 0\n",
    "    return jnp.logical_and(is_even, x > 0)  # Does not short-circuit, x > 0 is always evaluated\n",
    "\n",
    "\n",
    "print(python_check_positive_even(24))\n",
    "print(jax_check_positive_even(24))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7d001e3",
   "metadata": {},
   "source": [
    "Python logical operators cannot be applied element-wise:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "589418e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[False  True False]\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mValueError\u001b[39m                                Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[10]\u001b[39m\u001b[32m, line 3\u001b[39m\n\u001b[32m      1\u001b[39m x = jnp.array([\u001b[32m1\u001b[39m, \u001b[32m2\u001b[39m, \u001b[32m5\u001b[39m])\n\u001b[32m      2\u001b[39m \u001b[38;5;28mprint\u001b[39m(jax_check_positive_even(x))  \u001b[38;5;66;03m# Success! Can be applied element-wise\u001b[39;00m\n\u001b[32m----> \u001b[39m\u001b[32m3\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[43mpython_check_positive_even\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m)  \u001b[38;5;66;03m# Error! Cannot be applied element-wise\u001b[39;00m\n",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[9]\u001b[39m\u001b[32m, line 3\u001b[39m, in \u001b[36mpython_check_positive_even\u001b[39m\u001b[34m(x)\u001b[39m\n\u001b[32m      1\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mpython_check_positive_even\u001b[39m(x):\n\u001b[32m      2\u001b[39m     is_even = x % \u001b[32m2\u001b[39m == \u001b[32m0\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m3\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mis_even\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mand\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m      4\u001b[39m \u001b[43m        \u001b[49m\u001b[43mx\u001b[49m\u001b[43m \u001b[49m\u001b[43m>\u001b[49m\u001b[43m \u001b[49m\u001b[32;43m2\u001b[39;49m\n\u001b[32m      5\u001b[39m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/Programming/Learn JAX/.venv/lib/python3.12/site-packages/jax/_src/array.py:306\u001b[39m, in \u001b[36mArrayImpl.__bool__\u001b[39m\u001b[34m(self)\u001b[39m\n\u001b[32m    305\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34m__bool__\u001b[39m(\u001b[38;5;28mself\u001b[39m):\n\u001b[32m--> \u001b[39m\u001b[32m306\u001b[39m   \u001b[43mcore\u001b[49m\u001b[43m.\u001b[49m\u001b[43mcheck_bool_conversion\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[32m    307\u001b[39m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mbool\u001b[39m(\u001b[38;5;28mself\u001b[39m._value)\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/Programming/Learn JAX/.venv/lib/python3.12/site-packages/jax/_src/core.py:821\u001b[39m, in \u001b[36mcheck_bool_conversion\u001b[39m\u001b[34m(arr)\u001b[39m\n\u001b[32m    818\u001b[39m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[33m\"\u001b[39m\u001b[33mThe truth value of an empty array is ambiguous. Use\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m    819\u001b[39m                    \u001b[33m\"\u001b[39m\u001b[33m `array.size > 0` to check that an array is not empty.\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m    820\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m arr.size > \u001b[32m1\u001b[39m:\n\u001b[32m--> \u001b[39m\u001b[32m821\u001b[39m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[33m\"\u001b[39m\u001b[33mThe truth value of an array with more than one element\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m    822\u001b[39m                    \u001b[33m\"\u001b[39m\u001b[33m is ambiguous. Use a.any() or a.all()\u001b[39m\u001b[33m\"\u001b[39m)\n",
      "\u001b[31mValueError\u001b[39m: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()"
     ]
    }
   ],
   "source": [
    "x = jnp.array([1, 2, 5])\n",
    "print(jax_check_positive_even(x))  # Success! Can be applied element-wise\n",
    "print(python_check_positive_even(x))  # Error! Cannot be applied element-wise"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12ce2d0c",
   "metadata": {},
   "source": [
    "## Behavior in Grad\n",
    "\n",
    "Note that the Python control flow requirements are only in JIT contexts.\n",
    "\n",
    "Python control flow works fine in `@jax.grid` contexts without issues.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
